import torch


class SigDis(object):
    def __init__(self, aug_k_std, aug_b_std):
        self.aug = aug_k_std > 1e-6 and aug_b_std > 1e-6
        self.aug_k_std, self.aug_b_std = aug_k_std, aug_b_std
    
    def __call__(self, img_batch: torch.Tensor):
        if self.aug:
            k = torch.normal(mean=torch.ones_like(img_batch), std=self.aug_k_std)
            b = torch.normal(mean=torch.zeros_like(img_batch), std=self.aug_b_std)
            return img_batch * k + b
        else:
            return img_batch
